import _thread as thread
import time
import base64
import datetime
import hashlib
import hmac
import json
from typing import Optional, List, Dict
from urllib.parse import urlparse
import ssl
from datetime import datetime
from time import mktime
from urllib.parse import urlencode
from wsgiref.handlers import format_date_time

import websocket

# 用于存储 AI 回答的全局变量
# 注意：在多线程/异步环境中，这种全局变量需要更复杂的同步机制
# 对于简单的演示，可以接受，但生产环境应避免
global_ai_response = []
global_ws_closed = False # 标志 WebSocket 是否已关闭

class Ws_Param(object):
    """WebSocket连接参数类，用于生成和管理WebSocket连接所需的认证信息"""
    def __init__(self, APPID, APIKey, APISecret, gpt_url):
        self.APPID = APPID
        self.APIKey = APIKey
        self.APISecret = APISecret
        self.host = urlparse(gpt_url).netloc
        self.path = urlparse(gpt_url).path
        self.gpt_url = gpt_url

    def create_url(self):
        """生成带有认证信息的WebSocket连接URL"""
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 拼接签名原始字符串
        signature_origin = "host: " + self.host + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + self.path + " HTTP/1.1"

        # 使用hmac-sha256算法进行加密
        signature_sha = hmac.new(self.APISecret.encode('utf-8'), signature_origin.encode('utf-8'),
                                 digestmod=hashlib.sha256).digest()

        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')

        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'

        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')

        # 对加密结果进行base64编码
        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }

        url = self.gpt_url + '?' + urlencode(v)
        return url


# 处理WebSocket错误
def on_error(ws, error):
    print("### WebSocket Error:", error)
    global global_ai_response
    global_ai_response.append(f"Error: {error}")
    global global_ws_closed
    global_ws_closed = True


# 处理WebSocket关闭事件
def on_close(ws, close_status_code, close_msg):
    print("### WebSocket Closed ###")
    global global_ws_closed
    global_ws_closed = True


# 处理WebSocket连接建立事件
def on_open(ws):
    # 在新线程中运行发送数据逻辑，避免阻塞主线程
    thread.start_new_thread(run_websocket_send, (ws,))


def run_websocket_send(ws, *args):
    """
    在新线程中发送WebSocket数据
    """
    data = json.dumps(gen_params(appid=ws.appid, query=ws.query, history=ws.history, domain=ws.domain))
    ws.send(data)


# 处理WebSocket消息接收
def on_message(ws, message):
    global global_ai_response
    global global_ws_closed

    data = json.loads(message)
    code = data['header']['code']
    if code != 0:
        print(f'Request error: {code}, {data}')
        global_ai_response.append(f"Request error: {code}, {data.get('header', {}).get('message', 'Unknown error')}")
        ws.close()
    else:
        choices = data["payload"]["choices"]
        status = choices["status"]
        content = choices["text"][0]["content"]
        global_ai_response.append(content)

        if status == 2:
            print("#### 会话已结束 ####")
            ws.close()


def gen_params(appid, query, history: Optional[List[Dict]] = None, domain="generalv3.5"):
    """
    根据appid、用户查询和历史对话生成请求参数
    """
    message_text = []
    if history:
        for msg in history:
            message_text.append({"role": msg["role"], "content": msg["content"]})
    message_text.append({"role": "user", "content": query})

    data = {
        "header": {
            "app_id": appid,
            "uid": "1234", # 可根据实际需求使用动态用户ID
        },
        "parameter": {
            "chat": {
                "domain": domain,
                "temperature": 0.5,
                "max_tokens": 4096,
                "auditing": "default",
            }
        },
        "payload": {
            "message": {
                "text": message_text
            }
        }
    }
    return data


def get_ai_answer_from_spark(query: str, history: Optional[List[Dict]] = None) -> str:
    """
    调用讯飞星火AI助手获取回答
    这是一个同步封装函数，适用于在FastAPI服务中调用WebSocket客户端
    """
    global global_ai_response
    global global_ws_closed

    global_ai_response = []
    global_ws_closed = False

    APPID = "fc9821ef"
    API_SECRET = "YWZlOTA3MGU2ZDMzYzhiMjllODRmYTNj"
    API_KEY = "b07b7594a96a4f5886e24faae8c393ea"
    Spark_url = "wss://spark-api.xf-yun.com/v3.5/chat"
    DOMAIN = "generalv3.5"

    wsParam = Ws_Param(APPID, API_KEY, API_SECRET, Spark_url)
    websocket.enableTrace(False)
    wsUrl = wsParam.create_url()

    ws = websocket.WebSocketApp(wsUrl,
                                on_message=on_message,
                                on_error=on_error,
                                on_close=on_close,
                                on_open=on_open)
    ws.appid = APPID
    ws.query = query
    ws.history = history
    ws.domain = DOMAIN

    # 在单独线程中运行WebSocket客户端，避免run_forever阻塞主线程
    # 这能防止阻塞FastAPI主线程
    def run_ws():
        ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})

    ws_thread = thread.start_new_thread(run_ws, ())

    while not global_ws_closed:
        time.sleep(0.1)

    full_answer = "".join(global_ai_response)
    return full_answer

